#!/usr/bin/env python3
# 指定使用python3解释器
"""
MQTT-based prediction script for IoT sensor classification.
Subscribes to MQTT data, collects it into windows, and makes predictions.
"""
# 导入所需的库
import os  # 操作系统接口
import time  # 时间相关功能
import pandas as pd  # 数据处理库
import numpy as np  # 数值计算库
import joblib  # 模型保存/加载工具
import json  # JSON处理
import argparse  # 命令行参数解析
from typing import Dict, List, Any, Tuple, Optional  # 类型提示
from collections import deque  # 双端队列
from datetime import datetime  # 日期时间处理

import paho.mqtt.client as mqtt  # MQTT客户端库

from feature_extraction import extract_statistical_features  # 特征提取函数
from model import SensorClassifier  # 自定义分类器类


class MqttSensorDataCollector:
    """
    Class to collect sensor data from MQTT and maintain a buffer of recent data.
    """
    # MQTT传感器数据收集器类
    def __init__(self, broker_address: str, broker_port: int, topic: str, username: str = None, 
                 password: str = None, window_size: int = 100):
        """
        Initialize the MQTT sensor data collector.
        
        Args:
            broker_address: MQTT broker address
            broker_port: MQTT broker port
            topic: MQTT topic to subscribe to
            username: MQTT username for authentication (optional)
            password: MQTT password for authentication (optional)
            window_size: Size of the data window to maintain
        """
        # 初始化MQTT数据收集器
        self.broker_address = broker_address  # MQTT代理地址
        self.broker_port = broker_port  # MQTT代理端口
        self.topic = topic  # 订阅主题
        self.username = username  # 用户名(可选)
        self.password = password  # 密码(可选)
        self.window_size = window_size  # 数据窗口大小
        
        # 修改为按设备ID存储数据缓冲区的字典
        self.device_buffers = {}  # 格式: {device_id: deque()}
        
        # 修改为按设备ID存储状态的字典
        self.device_states = {}  # 格式: {device_id: {'has_enough_data': False}}
        
        # 初始化MQTT客户端
        self.client = mqtt.Client()
        self.client.on_connect = self.on_connect  # 连接回调
        self.client.on_message = self.on_message  # 消息回调
        
        # 如果提供了用户名和密码，设置认证
        if self.username and self.password:
            self.client.username_pw_set(self.username, self.password)
        
        # 是否有足够数据进行预测的标志
        self.has_enough_data = False
        
        # 最后接收到的设备ID
        self.device_id = None
    
    def on_connect(self, client, userdata, flags, rc):
        """Callback when connected to MQTT broker"""
        # MQTT连接回调函数
        connection_result = {
            0: "Connection successful",
            1: "Connection refused - incorrect protocol version",
            2: "Connection refused - invalid client identifier",
            3: "Connection refused - server unavailable",
            4: "Connection refused - bad username or password",
            5: "Connection refused - not authorized"
        }
        result_message = connection_result.get(rc, f"Unknown result code: {rc}")
        print(f"Connected to MQTT broker: {result_message}")
        
        if rc == 0:
            # 连接成功，订阅主题
            client.subscribe(self.topic)
            print(f"Subscribed to topic: {self.topic}")
        else:
            print(f"Failed to connect to MQTT broker. Check your credentials and connection settings.")
    
    def on_message(self, client, userdata, msg):
        """Callback when a message is received from MQTT broker"""
        try:
            # 解析JSON消息
            payload = json.loads(msg.payload.decode())
            
            # 从消息中提取数据
            device_id = payload.get('device')
            if not device_id:
                print("Warning: Message received without device ID")
                return
                
            timestamp = payload.get('ts')
            
            # 确保该设备的缓冲区存在
            if device_id not in self.device_buffers:
                self.device_buffers[device_id] = deque(maxlen=self.window_size)
                self.device_states[device_id] = {'has_enough_data': False}
            
            # 提取传感器数据
            acc = payload.get('acc', {})
            angle = payload.get('angle', {})
            
            # 创建记录
            record = {
                'acc_x': acc.get('x'),
                'acc_y': acc.get('y'),
                'acc_z': acc.get('z'),
                'angle_x': angle.get('x'),
                'angle_y': angle.get('y'),
                'angle_z': angle.get('z'),
                'timestamp': timestamp
            }
            
            # 将记录添加到对应设备的缓冲区
            self.device_buffers[device_id].append(record)
            
            # 更新设备数据状态
            if len(self.device_buffers[device_id]) >= self.window_size:
                self.device_states[device_id]['has_enough_data'] = True
            
        except Exception as e:
            print(f"Error processing MQTT message: {e}")
    
    def start(self):
        """Start the MQTT client"""
        # 启动MQTT客户端
        try:
            print(f"Connecting to MQTT broker at {self.broker_address}:{self.broker_port}")
            if self.username:
                print(f"Using authentication with username: {self.username}")
            
            self.client.connect(self.broker_address, self.broker_port, 60)
            # 以非阻塞方式启动MQTT客户端循环
            self.client.loop_start()
        except Exception as e:
            print(f"Error connecting to MQTT broker: {e}")
    
    def stop(self):
        """Stop the MQTT client"""
        # 停止MQTT客户端
        self.client.loop_stop()
        self.client.disconnect()
        print("Stopped MQTT client")
    
    def get_current_window(self, device_id: str) -> Optional[pd.DataFrame]:
        """
        Get the current window of sensor data for a specific device.
        
        Args:
            device_id: The ID of the device to get data for
            
        Returns:
            DataFrame containing the current window of sensor data or None if not enough data
        """
        if device_id not in self.device_states or not self.device_states[device_id]['has_enough_data']:
            return None
        
        # 将设备缓冲区转换为DataFrame
        df = pd.DataFrame(list(self.device_buffers[device_id]))
        
        # 确保所有必需的列都存在
        sensor_cols = ['acc_x', 'acc_y', 'acc_z', 'angle_x', 'angle_y', 'angle_z']
        for col in sensor_cols:
            if col not in df.columns:
                print(f"Warning: Column {col} not found in data for device {device_id}")
                df[col] = 0.0
        
        return df[sensor_cols]
    
    def get_active_devices(self) -> List[str]:
        """获取当前活跃的设备列表"""
        return list(self.device_buffers.keys())


def load_model_and_scaler(model_path: str, scaler_path: str) -> tuple:
    """
    Load the trained model and scaler.
    
    Args:
        model_path: Path to the trained model
        scaler_path: Path to the feature scaler
        
    Returns:
        Tuple of (classifier, scaler)
    """
    # 加载训练好的模型和特征缩放器
    classifier = SensorClassifier()  # 创建分类器实例
    classifier.load_model(model_path)  # 加载模型
    
    # 加载特征缩放器
    scaler = joblib.load(scaler_path)
    
    return classifier, scaler


def predict_with_mqtt(model_path: str, scaler_path: str, broker_address: str, broker_port: int, 
                     topic: str, username: str = None, password: str = None, 
                     window_size: int = 100, interval: float = 1.0):
    """
    Real-time prediction using data from MQTT.
    
    Args:
        model_path: Path to the trained model
        scaler_path: Path to the feature scaler
        broker_address: MQTT broker address
        broker_port: MQTT broker port
        topic: MQTT topic to subscribe to
        username: MQTT username for authentication (optional)
        password: MQTT password for authentication (optional)
        window_size: Size of the window in number of samples
        interval: Time interval between predictions in seconds
    """
    # 使用MQTT数据进行实时预测
    print(f"Loading model from {model_path} and scaler from {scaler_path}")
    classifier, scaler = load_model_and_scaler(model_path, scaler_path)
    
    # 创建MQTT数据收集器
    collector = MqttSensorDataCollector(
        broker_address, 
        broker_port, 
        topic, 
        username, 
        password, 
        window_size
    )
    
    # 为每个设备维护预测历史
    device_predictions = {}  # 格式: {device_id: [predictions]}
    
    collector.start()
    
    try:
        print(f"Starting prediction loop with interval {interval} seconds")
        while True:
            # 获取所有活跃设备
            active_devices = collector.get_active_devices()
            
            # 为每个活跃设备进行预测
            for device_id in active_devices:
                window = collector.get_current_window(device_id)
                
                if window is not None:
                    # 确保设备的预测历史存在
                    if device_id not in device_predictions:
                        device_predictions[device_id] = []
                    
                    # 提取特征并预测
                    features = extract_statistical_features(window)
                    features_df = pd.DataFrame([features])
                    normalized_features = pd.DataFrame(
                        scaler.transform(features_df),
                        columns=features_df.columns
                    )
                    
                    prediction = classifier.predict(normalized_features)[0]
                    prediction_proba = classifier.predict_proba(normalized_features)[0]
                    max_proba = max(prediction_proba)
                    
                    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    print(f"[{timestamp}] Device: {device_id}, Prediction: {prediction}, Confidence: {max_proba:.2f}")
                    
                    # 更新设备的预测历史
                    device_predictions[device_id].append(prediction)
                    if len(device_predictions[device_id]) > 5:
                        device_predictions[device_id] = device_predictions[device_id][-5:]
                        
                        # 获取最常见预测
                        prediction_counts = {}
                        for p in device_predictions[device_id]:
                            prediction_counts[p] = prediction_counts.get(p, 0) + 1
                        
                        most_common = max(prediction_counts, key=prediction_counts.get)
                        print(f"Device {device_id} most common prediction (last 5): {most_common}")
                        
                        # 发布预测结果
                        mqtt_payload = json.dumps({
                            "device": device_id,
                            "posture": most_common.upper(),
                            "confidence": float(max_proba)
                        })
                        if topic.find("DeskData") != -1:
                            collector.client.publish(f"Demo/AIIdentify/{device_id}", mqtt_payload)
                        else:
                            collector.client.publish(f"Demo/Posture/Predict/{device_id}", mqtt_payload)
            
            time.sleep(interval)
    
    except KeyboardInterrupt:
        print("Prediction stopped by user")
    finally:
        # 停止收集器
        collector.stop()


if __name__ == "__main__":
    # 获取项目根目录的绝对路径
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="MQTT-based prediction for IoT sensor classification")
    
    parser.add_argument("--model", type=str, default=os.path.join(project_root, "output", "model.joblib"),
                        help="Path to the trained model")
    parser.add_argument("--scaler", type=str, default=os.path.join(project_root, "output", "scaler.joblib"),
                        help="Path to the feature scaler")
    parser.add_argument("--broker", type=str, required=True,
                        help="MQTT broker address")
    parser.add_argument("--port", type=int, default=1883,
                        help="MQTT broker port")
    parser.add_argument("--topic", type=str, default="sensor/data",
                        help="MQTT topic to subscribe to")
    parser.add_argument("--username", type=str, default=None,
                        help="MQTT username for authentication")
    parser.add_argument("--password", type=str, default=None,
                        help="MQTT password for authentication")
    parser.add_argument("--window-size", type=int, default=100,
                        help="Size of the window in number of samples")
    parser.add_argument("--interval", type=float, default=1.0,
                        help="Time interval between predictions in seconds")
    
    args = parser.parse_args()
    
    # 开始预测
    predict_with_mqtt(
        model_path=args.model,
        scaler_path=args.scaler,
        broker_address=args.broker,
        broker_port=args.port,
        topic=args.topic,
        username=args.username,
        password=args.password,
        window_size=args.window_size,
        interval=args.interval
    )
